import os
import json
from PIL import Image


# 读取YOLO格式的标注文件
def read_yolo_file(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()
    return lines

# 解析YOLO格式数据为JSON格式
def parse_yolo_to_json(yolo_data, image_path, image_height, image_width):
    json_data = {
        "version": "5.5.0",
        "flags": {},
        "shapes": [],
        "imagePath": image_path,
        "imageData": None,
        "imageHeight": image_height,
        "imageWidth": image_width
    }

    for line in yolo_data:
        line_data = line.strip().split()
        label = ["box", "arrows_right", "arrows_left"][int(line_data[0])]
        center_x, center_y, width, height = map(float, line_data[1:5])
        # 反归一化边界框坐标
        center_x = center_x * image_width
        center_y = center_y * image_height
        width = width * image_width
        height = height * image_height
        points = [[center_x - width/2, center_y - height/2], [center_x + width/2, center_y + height/2]]
        shape = {
            "label": label,
            "points": points,
            "group_id": None,
            "description": "",
            "shape_type": "rectangle",
            "flags": {},
            "mask": None
        }
        json_data["shapes"].append(shape)

        for i in range(5, len(line_data), 3):
            x, y, visible = map(float, line_data[i:i+3])
            x = x * image_width
            y = y * image_height
            save_label = f"{label}_{i//3}"
            shape = {
                "label": save_label,
                "points": [[x, y]],
                "group_id": None,
                "description": "",
                "shape_type": "point",
                "flags": {},
                "mask": None
            }
            json_data["shapes"].append(shape)

    return json_data

# 保存JSON数据到文件
def save_json_to_file(json_data, file_path):
    with open(file_path, 'w') as file:
        json.dump(json_data, file, indent=4)


# 文件夹A中的路径
folder_a_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/good/仅正面/yolo_annotations"
output_folder = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/good/仅正面/json_annotations"


os.makedirs(output_folder, exist_ok=True)

# 遍历文件夹A中的文件
for root, dirs, files in os.walk(folder_a_path):
    for file_name in files:
        if file_name.endswith(".jpg"):
            image_file_path = os.path.join(root, file_name)
            txt_file_path = os.path.splitext(image_file_path)[0] + ".txt"

            if os.path.exists(txt_file_path):
                # 读取YOLO标注数据
                yolo_data = read_yolo_file(txt_file_path)

                if len(yolo_data) == 0:
                    continue

                # 获取图片信息
                image = Image.open(image_file_path)
                image_width, image_height = image.size

                # 解析YOLO数据为JSON格式
                json_data = parse_yolo_to_json(yolo_data, file_name, image_height, image_width)

                # 保存JSON数据到文件夹B
                json_file_path = os.path.join(output_folder, os.path.splitext(file_name)[0] + ".json")
                save_json_to_file(json_data, json_file_path)

print("转换完成并保存到文件夹B中。")